/*
 * accerelated_gradient.h
 *
 *  Created on: Jun 3, 2011
 *      Author: tqlong
 */

#ifndef ACCERELATED_GRADIENT_H_
#define ACCERELATED_GRADIENT_H_

#include "line_search.h"

template<typename Function>
  class AcceleratedGradient
  {
  public:
    typedef Function func_type;
    typedef LineSearch<Function> linesearch_type;
    typedef typename Function::range_type range_type;
    typedef typename Function::domain_type domain_type;
  public:
    AcceleratedGradient(func_type &f, const domain_type &x0, domain_type &x) :
      f_(f), x0_(x0), y_(&x), step_(1.0)
    {
      if (!f.has_value_)
        {
          std::cerr << "Function needs to produce value() at " << __FILE__
              << ":" << __LINE__ << std::endl;
          exit(1);
        }
      if (!f.has_gradient_)
        {
          std::cerr << "Function needs to produce gradient() at " << __FILE__
              << ":" << __LINE__ << std::endl;
          exit(1);
        }
    }

    bool
    search()
    {
      yprev_ = *y_ = x0_;
      tprev_ = 0;
      t_ = 1;
      iter_ = 0;
      for (iter_ = 0; iter_ < 100; iter_++)
        {
          double alpha = (tprev_ - 1) / t_;
          x_ = (1 + alpha) * (*y_) - alpha * yprev_;

          f_.value_gradient(x_, fx_, grad_);
          domain_type dir = -grad_;

          yprev_ = *y_;
          linesearch_type ls(f_, x_, *y_, fx_, grad_, dir, step_);

          bool ok = ls.search();
          step_ = ls.step();
          tprev_ = t_;
          t_ = 0.5 * (1 + sqrt(1 + 4 * tprev_ * tprev_));

          double grad_norm = arma::norm(grad_, 2);

          std::cout << "iter=" << iter_ << " fx=" << ls.fx() << " step=" << step_
              << " |grad|=" << grad_norm << std::endl;
          if (!ok || grad_norm < 1e-5) break; // line search failed
        }
      f_.value_gradient(*y_, fx_, grad_);
      return true;
    }

    void
    value_gradient(const domain_type& x, range_type& value, domain_type& grad)
    {
      if (f_.has_both_)
        f_.value_gradient(x, value, grad);
      else
        {
          value = f_.value(x);
          grad = f_.gradient(x);
        }
    }

    range_type
    fx() const
    {
      return fx_;
    }

    const domain_type&
    grad() const
    {
      return grad_;
    }
  protected:
    func_type &f_;
    const domain_type &x0_;
    domain_type yprev_, *y_, x_;
    range_type fx_;
    domain_type grad_;
    double step_, t_, tprev_;
    int iter_;
  };

#endif /* ACCERELATED_GRADIENT_H_ */
